"""
Docker容器实时日志WebSocket Consumer
"""
import json
import asyncio
import logging
import threading
import queue
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from docker_management.models import DockerServer
import paramiko

logger = logging.getLogger(__name__)


class DockerLogsConsumer(AsyncWebsocketConsumer):
    """Docker容器实时日志WebSocket Consumer"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_client = None
        self.ssh_channel = None
        self.server_id = None
        self.container_id = None
        self.docker_server = None
        self.user = None
        self.is_connected = False
        self.read_thread = None
        self.output_queue = queue.Queue()
        self.output_task = None
        self.follow_logs = True

    async def connect(self):
        """WebSocket连接建立"""
        try:
            # 获取用户信息（从中间件设置）
            self.user = self.scope.get("user")

            if not self.user or not getattr(self.user, 'is_authenticated', False):
                logger.warning("Docker日志WebSocket连接被拒绝：用户未认证")
                await self.close(code=4001)
                return

            # 获取服务器ID和容器ID
            self.server_id = self.scope['url_route']['kwargs']['server_id']
            self.container_id = self.scope['url_route']['kwargs']['container_id']

            logger.info(f"Docker日志连接请求: 用户={self.user.username}, 服务器={self.server_id}, 容器={self.container_id}")

            # 接受WebSocket连接
            await self.accept()

            # 发送连接状态
            await self.send(text_data=json.dumps({
                'type': 'status',
                'status': 'websocket_connected',
                'message': 'WebSocket连接已建立'
            }))

        except Exception as e:
            logger.error(f"Docker日志WebSocket连接失败: {e}")
            await self.close(code=4000)

    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        logger.info(f"Docker日志WebSocket连接断开: {close_code}")
        await self.cleanup_connection()

    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')

            if message_type == 'start_logs':
                # 开始获取日志
                lines = data.get('lines', 100)
                follow = data.get('follow', True)
                await self.start_logs(lines=lines, follow=follow)
            elif message_type == 'stop_logs':
                # 停止获取日志
                await self.stop_logs()
            else:
                logger.warning(f"未知的消息类型: {message_type}")

        except json.JSONDecodeError:
            logger.error("无效的JSON消息")
        except Exception as e:
            logger.error(f"处理WebSocket消息失败: {e}")

    async def start_logs(self, lines=100, follow=True):
        """开始获取容器日志"""
        try:
            # 获取Docker服务器信息
            self.docker_server = await database_sync_to_async(
                DockerServer.objects.get
            )(pk=self.server_id, is_active=True)

            # 建立SSH连接
            await self.connect_ssh()

            # 设置日志参数
            self.follow_logs = follow

            # 开始获取日志
            await self.exec_docker_logs(lines=lines, follow=follow)

        except DockerServer.DoesNotExist:
            await self.send_error("Docker服务器不存在")
        except Exception as e:
            logger.error(f"开始获取日志失败: {e}")
            await self.send_error(f"开始获取日志失败: {str(e)}")

    async def stop_logs(self):
        """停止获取日志"""
        try:
            self.follow_logs = False
            self.is_connected = False

            await self.send(text_data=json.dumps({
                'type': 'status',
                'status': 'logs_stopped',
                'message': '日志获取已停止'
            }))

        except Exception as e:
            logger.error(f"停止日志失败: {e}")

    async def connect_ssh(self):
        """建立SSH连接"""
        try:
            def create_ssh_connection():
                ssh_client = paramiko.SSHClient()
                ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                # 连接SSH
                ssh_client.connect(
                    hostname=self.docker_server.host,
                    port=self.docker_server.ssh_port,
                    username=self.docker_server.ssh_username,
                    password=self.docker_server.ssh_password,
                    timeout=10
                )

                return ssh_client

            # 在线程池中执行SSH连接
            loop = asyncio.get_event_loop()
            self.ssh_client = await loop.run_in_executor(None, create_ssh_connection)

            logger.info(f"SSH连接成功: {self.docker_server.name}")

        except Exception as e:
            logger.error(f"SSH连接失败: {e}")
            raise

    async def exec_docker_logs(self, lines=100, follow=True):
        """执行docker logs命令"""
        try:
            def execute_logs():
                try:
                    # 构建docker logs命令
                    cmd = f'docker logs --tail {lines}'
                    if follow:
                        cmd += ' -f'
                    cmd += f' {self.container_id}'

                    logger.info(f"执行命令: {cmd}")

                    # 执行命令
                    stdin, stdout, stderr = self.ssh_client.exec_command(cmd)

                    return stdout, stderr

                except Exception as e:
                    logger.error(f"执行docker logs失败: {e}")
                    raise e

            # 在线程池中执行
            loop = asyncio.get_event_loop()
            stdout, stderr = await loop.run_in_executor(None, execute_logs)

            self.is_connected = True

            # 发送开始消息
            await self.send(text_data=json.dumps({
                'type': 'logs_started',
                'message': f'开始获取容器 {self.container_id} 的日志'
            }))

            # 启动读取线程
            self.start_read_thread(stdout, stderr)

        except Exception as e:
            logger.error(f"执行docker logs失败: {e}")
            raise

    def start_read_thread(self, stdout, stderr):
        """启动读取线程"""
        def read_from_streams():
            try:
                while self.is_connected and self.follow_logs:
                    # 读取stdout
                    if stdout.channel.recv_ready():
                        data = stdout.channel.recv(4096)
                        if data:
                            self.output_queue.put({
                                'type': 'stdout',
                                'data': data.decode('utf-8', errors='ignore')
                            })

                    # 读取stderr
                    if stderr.channel.recv_ready():
                        data = stderr.channel.recv(4096)
                        if data:
                            self.output_queue.put({
                                'type': 'stderr',
                                'data': data.decode('utf-8', errors='ignore')
                            })

                    # 检查是否还有数据
                    if stdout.channel.exit_status_ready():
                        # 读取剩余数据
                        remaining_stdout = stdout.read().decode('utf-8', errors='ignore')
                        remaining_stderr = stderr.read().decode('utf-8', errors='ignore')

                        if remaining_stdout:
                            self.output_queue.put({
                                'type': 'stdout',
                                'data': remaining_stdout
                            })
                        if remaining_stderr:
                            self.output_queue.put({
                                'type': 'stderr',
                                'data': remaining_stderr
                            })
                        break

                    import time
                    time.sleep(0.1)

            except Exception as e:
                logger.error(f"读取日志输出失败: {e}")

        self.read_thread = threading.Thread(target=read_from_streams, daemon=True)
        self.read_thread.start()

        # 启动输出处理任务
        self.output_task = asyncio.create_task(self.process_output_queue())

    async def process_output_queue(self):
        """处理输出队列"""
        try:
            while self.is_connected:
                try:
                    # 非阻塞地获取队列中的数据
                    log_data = self.output_queue.get_nowait()
                    await self.send_log_data(log_data)
                except queue.Empty:
                    # 队列为空，等待一下
                    await asyncio.sleep(0.01)
                except Exception as e:
                    logger.error(f"处理输出队列失败: {e}")
                    break
        except Exception as e:
            logger.error(f"输出队列处理任务失败: {e}")

    async def send_log_data(self, log_data):
        """发送日志数据到前端"""
        await self.send(text_data=json.dumps({
            'type': 'log_data',
            'stream': log_data['type'],  # stdout 或 stderr
            'data': log_data['data']
        }))

    async def send_error(self, message):
        """发送错误消息"""
        await self.send(text_data=json.dumps({
            'type': 'error',
            'message': message
        }))

    async def cleanup_connection(self):
        """清理连接"""
        try:
            self.is_connected = False
            self.follow_logs = False

            # 取消输出处理任务
            if self.output_task and not self.output_task.done():
                self.output_task.cancel()
                try:
                    await self.output_task
                except asyncio.CancelledError:
                    pass

            if self.ssh_client:
                self.ssh_client.close()
                self.ssh_client = None

            if self.read_thread and self.read_thread.is_alive():
                self.read_thread.join(timeout=1)

            logger.info("Docker日志连接已清理")

        except Exception as e:
            logger.error(f"清理连接失败: {e}")
